package org.zjvis.datascience.common.util.sqlParse;

import org.apache.poi.ss.usermodel.Cell;
import org.apache.poi.ss.usermodel.Row;
import org.apache.poi.ss.usermodel.Sheet;
import org.apache.poi.ss.usermodel.Workbook;
import org.apache.poi.xssf.usermodel.XSSFWorkbook;
import org.zjvis.datascience.common.dto.DataDto;

import java.io.BufferedWriter;
import java.io.IOException;
import java.io.OutputStream;
import java.io.OutputStreamWriter;
import java.nio.charset.StandardCharsets;
import java.util.Map;
import java.util.Set;

/**
 * @description : 数据转文件帮助类
 * @date 2021-10-25
 */
public class DataToFileUtil {
    public static void DataToCsvStream(DataDto data, OutputStream outputStream) throws IOException {
        BufferedWriter bufferedWriter = new BufferedWriter(new OutputStreamWriter(outputStream, StandardCharsets.UTF_8));
        Set<String> headKey = data.getHead().keySet();
        String head = headKey.toString();
        bufferedWriter.write(head.substring(1, head.length() - 1));
        bufferedWriter.newLine();
        for (Map<String, String> column : data.getData()) {
            StringBuffer stringBuffer = new StringBuffer();
            for (String key : headKey) {
                if (column.get(key) != null) {
                    String value = column.get(key);
                    if (value.contains(",")) {
                        stringBuffer.append("\"").append(value).append("\"");
                    } else {
                        stringBuffer.append(value);
                    }
                }
                stringBuffer.append(",");
            }
            stringBuffer.deleteCharAt(stringBuffer.length() - 1);
            bufferedWriter.write(stringBuffer.toString());
            bufferedWriter.newLine();
        }
        bufferedWriter.flush();
        bufferedWriter.close();
    }

    public static Workbook DataToXlsx(DataDto data) {
        XSSFWorkbook workbook = new XSSFWorkbook();
        Sheet sheet = workbook.createSheet();
        Map<String, String> head = data.getHead();
        Row headRow = sheet.createRow(0);
        for (String column : head.keySet()) {
            Cell cell = headRow.createCell(headRow.getLastCellNum() == -1 ? headRow.getLastCellNum() + 1 : headRow.getLastCellNum());
            cell.setCellValue(column);
        }
        for (Map<String, String> rows : data.getData()) {
            Row row = sheet.createRow(sheet.getLastRowNum() + 1);
            for (String column : head.keySet()) {
                Cell cell = row.createCell(row.getLastCellNum() == -1 ? row.getLastCellNum() + 1 : row.getLastCellNum());
                cell.setCellValue(rows.get(column));
            }
        }
        return workbook;
    }
}
